#include "hllc.hpp"
#include "util.hpp"
#include "unifiedMathFunctions.hpp"
#include "sweFunctions.hpp"

namespace hms::hllc {

/** @brief 
 * 
 * Implements Simons (2019) Eq. 3.56 */

/**
 * @brief Returns the water depth in the star region.
 * 
 * @param vNormalL Velocity normal to the edge, left side.
 * @param vNormalR Velocity normal to the edge, right side.
 * @param cwL The wavespeed c_WL, not the HLLC one.
 * @param cwR The wavespeed c_WR, not the HLLC one.
 * 
 * Implements Simons (2019) Eq. 3.56. Does not work with flipped sides. */
scalar dStar(
	scalar vNormalL, scalar vNormalR, scalar cwL, scalar cwR
){
	return square(
		( cwL + cwR ) / 2 +
		( vNormalL - vNormalR ) / 4
	) / phys::g;
}


/** @brief Returns the middle wave speed (star region).
 * 
 * @param dL Reconstructed water depth, left side.
 * @param dR Reconstructed water depth, right side.
 * @param vNormalL Velocity normal to the edge, left side.
 * @param vNormalR Velocity normal to the edge, right side.
 * @param cwL The HLLC wavespeed, left side.
 * @param cwR The HLLC wavespeed, right side.
 * 
 * Implements Simons (2019), Eq. 3.58. */
scalar wavespeedStar(
	scalar dL, scalar dR,
	scalar vNormalL, scalar vNormalR,
	scalar SL, scalar SR
){
	assert( dL >= 0 && dR >= 0 &&
		"Water depths must not be negative!" );

	scalar helperL { dL * ( vNormalL - SL ) };
	scalar helperR { dR * ( vNormalR - SR ) };

	assert( helperL != helperR && "Division by zero detected!" );

	return
		( SL * helperR - SR * helperL ) / ( helperR - helperL );
}

/* alternatively, instead of calculating wavespeedStar: */
bool useFluxLeft(
	scalar dL, scalar dR,
	scalar uL, scalar uR,
	scalar SL, scalar SR
){
    scalar hL { dL * ( uL - SL ) };
    scalar hR { dR * ( uR - SR ) };
    return std::signbit( SL*hR - SR*hL ) == std::signbit( hR - hL );
}


/** @brief Returns the three components of the flux vector. Use when the
 * HLLC waves are all in one quadrant of the x-t-chart.
 * 
 * @param dSide The reconstructed water depth on the flux' side.
 * @param vNormalSide The normal velocity on the flux' side.
 * @param vSide The un-transformed flow velocity on the flux' side.
 * @param normal The normal vector of the edge.
 * @return Array3s 
 * 
 * Implements Simons (2019), Eq. 3.63 in an equivalent form (retransformation):
 * vSideX = normalX * vNormal - normalY * vParallel
 * vSideY = normalY * vNormal + normalX * vParallel
 */
Array3s flux(
	scalar dSide,
	scalar vNormalSide,
	const Ref<const Array2s>& vSide,
	const Ref<const Vector2s>& normal
){
	assert( dSide >= 0 && "Water depth must not be negative!" );
	assert( approxEqual( normal.norm(), 1 ) &&
		"Normal vector is not a unit vector!" );
	Array3s f;
	f = vNormalSide * dSide; // these are part of all components.
	f.tail(2) = f.tail(2) * vSide + normal.array() * pressureFlux(dSide);
	return f;
}


/**
 * @brief Returns the three components of the flux vector in the star region.
 * 
 * @param dSide The reconstructed water depth on the flux' side.
 * @param dOtherSide The reconstructed water depth on the other side.
 * @param vRotSide The flow velocity in the rotated coordinate system of the
 * edge's normal vector, on the flux' side.
 * @param vRotOtherSide The rotated flow velocity on the other side.
 * @param wavespeedSide The HLLC wave speed (S_side) on the flux' side.
 * @param wavespeedOtherSide The HLLC wave speed on the other side.
 * @param normal The edge's normal vector.
 * @return Array3s 
 * 
 * Implements Simons (2019), Eq. 3.64 with 3.53, which yields 3.65
 */
Array3s fluxStar(
	scalar dSide,
	scalar dOtherSide,
	const Ref<const Array2s>& vRotSide,
	const Ref<const Array2s>& vRotOtherSide,
	scalar wavespeedSide,
	scalar wavespeedOtherSide,
	const Ref<const Vector2s>& normal
){
	assert( dSide >= 0 && dOtherSide >= 0 &&
		"Water depths must not be negative!" );
	assert( approxEqual( normal.norm(), 1 ) &&
		"Normal vector is not a unit vector!" );

	scalar massFlux = hllFlux(
		wavespeedSide, wavespeedOtherSide,
		vRotSide.x() * dSide, vRotOtherSide.x() * dOtherSide,
		dSide, dOtherSide
	);

	scalar momentumFlux = hllFlux(
		wavespeedSide, wavespeedOtherSide,
		square( vRotSide.x() ) * dSide + pressureFlux(dSide),
		square( vRotOtherSide.x() ) * dOtherSide + pressureFlux(dOtherSide),
		vRotSide.x() * dSide, vRotOtherSide.x() * dOtherSide
	);

	return {
		massFlux,
		normal.x() * momentumFlux - normal.y() * vRotSide.y() * massFlux,
		normal.y() * momentumFlux + normal.x() * vRotSide.y() * massFlux,
	};
}


Array3s computeFlux(
	const EdgeValues& ev,
	const ArrayNNs<2,2>& vRotLR,
	const Vector2s& normal
){
	/* wave speeds */
	RowArray2s cwLR {
		ev.dLR.unaryExpr( [](scalar d){ return hms::wavespeed(d); } )
	};

	/* water depth and wave speed in the star region */
	scalar vStar { hllc::vNormalStar(
		vRotLR(0, 0), vRotLR(0, 1), cwLR[0], cwLR[1]
	)};
	scalar cwStar { hllc::cwStar(
		vRotLR(0, 0), vRotLR(0, 1), cwLR[0], cwLR[1]
	)};

	scalar SL, SR;

	using s = hllc::WaveRegion;

	SL = (ev.dLR[0] == 0) ?
		hllc::wavespeed_dry<s::left>( vRotLR(0,1), cwLR[1] ) :
		hllc::wavespeed<s::left>( vRotLR(0,0), vStar, cwLR[0], cwStar );
	SR = (ev.dLR[1] == 0) ?
		hllc::wavespeed_dry<s::right>( vRotLR(0,0), cwLR[0] ) :
		hllc::wavespeed<s::right>( vRotLR(0,1), vStar, cwLR[1], cwStar );

	Array3s flux;
	/* most likely case first */
	if (SL < 0 && SR > 0){
		// scalar SStar { hllc::wavespeedStar(
		// 	dLR[0], dLR[1], vRotLR(0,0), vRotLR(0,1), SL, SR
		// )};
		// if ( SStar >= 0)
		/* actually calculating SStar is not needed. What's needed is to
		 * know whether it's positive or negative */
		if ( hllc::useFluxLeft(
			ev.dLR[0], ev.dLR[1], vRotLR(0,0), vRotLR(0,1), SL, SR
		) )
			flux = hllc::fluxStar(
				ev.dLR[0], ev.dLR[1], vRotLR.col(0), vRotLR.col(1),
				SL, SR, normal
			);
		else
			flux = hllc::fluxStar(
				ev.dLR[1], ev.dLR[0], vRotLR.col(1), vRotLR.col(0),
				SR, SL, normal
			);
	} 
	
	else if (SL >= 0)
		flux = hllc::flux( ev.dLR[0], vRotLR(0,0), ev.vLR.col(0), normal );
	else
		flux = hllc::flux( ev.dLR[1], vRotLR(0,1), ev.vLR.col(1), normal );

	return flux;
}


} // end namespace hms::hllc
